module Data.Fits.FitsParser where

import qualified Data.ByteString.Lazy as B
import Data.Binary.Get
import Data.ByteString.Internal (w2c)
import Data.Fits.KeywordRecord
import Data.Fits.Fits
import Data.Fits.RecordParser
import Control.Monad.Error

type ErrorG = ErrorT String Get

readFits :: FilePath -> IO [Either String HDU]
readFits f = do
    b <- B.readFile f
    return $ parseFits b

parseFits :: B.ByteString -> [Either String HDU]
parseFits = runGet $ do x <- runErrorT fitsParser
                        case x of
                          Left  e  -> return [Left e]
                          Right x  -> return x

fitsParser :: ErrorG [Either String HDU]
fitsParser = do
    p <- loadPrimaryHDU
    o <- loadExtensionHDU `catchError` \e -> return [Left e]
    return (p:o)    

loadPrimaryHDU = do
    h <- loadHeader
    d <- loadPrimaryData h
    return $ Right $ HDU {header = h, unit = d}

loadExtensionHDU = do
    r <- lift remaining
    if r == 0 
      then return []
      else do h <- loadHeader
              d <- loadExtensionData h 
              s <- loadExtensionHDU `catchError` \e -> return [Left e]
              let hdu = Right $ HDU {header = h, unit = d}
              return (hdu:s)

loadHeader = do
  xs <- loadKeywords
  forM xs $ \x -> do
     let s = (map (w2c) . B.unpack) x
     case (parseRecord s) of
       Left e  -> let message = "Error in line\n" ++ "{" ++ s ++ "}" ++ "\n" ++ show e
                  in throwError $ strMsg message
       Right a -> return a     

loadKeywords = do
  block <- lift $ getLazyByteString 2880
  guard (B.length block == 2880) `catchError` \_ -> throwError $ strMsg "Size of Keyword block not divide by 2880"
  let (x,e) = lookBlock block
  if e 
    then return x
    else loadKeywords >>= \xs -> return $ x ++ xs
                           
lookBlock b = runGet lookBlockParser b
lookBlockParser = do
  x <- remaining
  if x == 0 
    then return ([], False)
    else do e <- lookAhead end
            if e 
              then return ([],True)
              else do l <- getLazyByteString 80
                      (ls,e) <- lookBlockParser
                      return (l:ls,e)

end = do
    x <- getLazyByteString 4
    let y = map (w2c) $ B.unpack x
    return $ y == "END "

attachExc m e = do case m of
                     Nothing -> throwError e
                     Just x  -> return x

round2880 x = if r == 0 
                then x 
                else 2880 + x - r 
   where r = x `mod` 2880

simpleHDUBits h = do
  bitpix  <- lookValue h "BITPIX"
  guard (bitpix `mod` 8 == 0)
  naxis   <- lookValue h "NAXIS" :: Maybe Integer
  naxises <- sequence $ map (lookValue h) ["NAXIS" ++ show x | x <- [1..naxis]]
  return $ (abs bitpix * product naxises) `div` 8

extensionHDUBits h = do
  bitpix  <- lookValue h "BITPIX"
  guard (bitpix `mod` 8 == 0)
  naxis   <- lookValue h "NAXIS" :: Maybe Integer
  naxises <- sequence $ map (lookValue h) ["NAXIS" ++ show x | x <- [1..naxis]]
  pcount  <- lookValue h "PCOUNT"
  gcount  <- lookValue h "GCOUNT"
  return $ (abs bitpix * gcount * (pcount + product naxises)) `div` 8

readChunk s = do
  let r = round2880 s - s
  x <- lift $ getLazyByteString $ fromInteger s
  lift $ getLazyByteString $ fromInteger r
  let len = B.length x
  guard (len == fromInteger s) `catchError` \_ -> throwError $ 
    strMsg $ "Can't read data block\n remaining: " ++ show len ++ "\n size:" ++ show s
  return x

loadPrimaryData h = do
  do { lookValue h "SIMPLE" >>= guard } 
       `attachExc` strMsg "Primary HDU must have SIMPLE key with T"
  s <- simpleHDUBits h `attachExc` strMsg "Can't determinate size of data block"
  readChunk s

loadExtensionData h = do
  (lookValue h "XTENSION" :: Maybe String) 
       `attachExc` strMsg "EXTENSION HDU must have keyword XTENSION"
  s <- extensionHDUBits h `attachExc` strMsg "Can't determinate size of data block"
  readChunk s
